#include <mcp/server/tool_manager.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <iostream>
#include <sstream>
#include <algorithm>
#include <memory>

namespace mcp {
namespace server {

// ToolManager implementation
struct ToolManager::Impl {
    std::unordered_map<std::string, ToolRegistration> tools;
    std::vector<std::string> tool_order; // For maintaining insertion order

    // Thread safety
    mutable std::mutex mutex;
};

ToolManager::ToolManager() : impl_(std::make_unique<Impl>()) {}

ToolManager::~ToolManager() = default;

bool ToolManager::register_tool(ToolRegistration registration) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    if (impl_->tools.contains(registration.name)) {
        return false; // Tool already exists
    }

    impl_->tools[registration.name] = std::move(registration);
    impl_->tool_order.push_back(registration.name);
    return true;
}

bool ToolManager::unregister_tool(const std::string& name) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->tools.find(name);
    if (it == impl_->tools.end()) {
        return false; // Tool not found
    }

    impl_->tools.erase(it);
    impl_->tool_order.erase(
        std::remove(impl_->tool_order.begin(), impl_->tool_order.end(), name),
        impl_->tool_order.end()
    );
    return true;
}

bool ToolManager::has_tool(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    return impl_->tools.contains(name);
}

std::vector<types::Tool> ToolManager::list_tools() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    std::vector<types::Tool> result;
    result.reserve(impl_->tool_order.size());

    for (const auto& name : impl_->tool_order) {
        auto it = impl_->tools.find(name);
        if (it != impl_->tools.end()) {
            const auto& registration = it->second;
            result.push_back(types::Tool{
                .name = registration.name,
                .description = registration.description,
                .input_schema = registration.input_schema,
                .output_schema = registration.output_schema
            });
        }
    }

    return result;
}

std::optional<types::Tool> ToolManager::get_tool(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->tools.find(name);
    if (it == impl_->tools.end()) {
        return std::nullopt;
    }

    const auto& registration = it->second;
    return types::Tool{
        .name = registration.name,
        .description = registration.description,
        .input_schema = registration.input_schema,
        .output_schema = registration.output_schema
    };
}

Result<ToolResult> ToolManager::call_tool(const std::string& name,
                                          const json::Value& arguments,
                                          Context<>& context) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->tools.find(name);
    if (it == impl_->tools.end()) {
        return Error::method_not_found(name);
    }

    const auto& registration = it->second;

    try {
        // Validate arguments against schema
        auto validation = registration.input_schema.validate(arguments);
        if (!validation.valid) {
            std::string error_msg = "Invalid arguments for tool '" + name + "':\n";
            for (const auto& error : validation.errors) {
                error_msg += "  - " + error + "\n";
            }
            return Error::invalid_params(error_msg);
        }

        // Call the tool handler
        auto result = registration.handler(arguments, context);
        if (!result) {
            return result.error();
        }

        // Convert the result to ToolResult format
        ToolResult tool_result;
        if (result->is_object() && result->contains("content")) {
            // If the result already has content structure, use it directly
            const auto& content_array = (*result)["content"];
            if (content_array.is_array()) {
                for (const auto& content_item : content_array.as_array()) {
                    if (content_item.is_object() && content_item.contains("type")) {
                        std::string type_str = content_item["type"].as_string();
                        if (type_str == "text" && content_item.contains("text")) {
                            tool_result.content.push_back(
                                std::make_shared<types::TextContent>(content_item["text"].as_string())
                            );
                        } else if (type_str == "image" && content_item.contains("data") &&
                                  content_item.contains("mime_type")) {
                            tool_result.content.push_back(
                                std::make_shared<types::ImageContent>(
                                    content_item["data"].as_string(),
                                    content_item["mime_type"].as_string()
                                )
                            );
                        }
                    }
                }
            }
        }

        // If no structured content was found, create a text content from the result
        if (tool_result.content.empty()) {
            tool_result.content.push_back(
                std::make_shared<types::TextContent>(result->serialize())
            );
        }

        // Check if the result indicates an error
        if (result->is_object() && result->contains("isError")) {
            tool_result.is_error = result->get("isError", false);
        }

        // Include structured content if present
        if (result->is_object() && result->contains("structured_content")) {
            tool_result.structured_content = result->get("structured_content", json::Value());
        }

        return tool_result;

    } catch (const std::exception& e) {
        return Error::internal_error(
            std::string("Exception in tool '") + name + "': " + e.what()
        );
    } catch (...) {
        return Error::internal_error(
            std::string("Unknown exception in tool '") + name + "'"
        );
    }
}

json::ValidationResult ToolManager::validate_arguments(const std::string& name,
                                                      const json::Value& arguments) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->tools.find(name);
    if (it == impl_->tools.end()) {
        return json::ValidationResult{false, {"Tool not found: " + name}};
    }

    return it->second.input_schema.validate(arguments);
}

std::optional<types::ToolAnnotations> ToolManager::get_annotations(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->tools.find(name);
    if (it == impl_->tools.end()) {
        return std::nullopt;
    }

    return it->second.annotations;
}

void ToolManager::clear() {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    impl_->tools.clear();
    impl_->tool_order.clear();
}

// ToolResult utilities
ToolResult create_text_result(const std::string& text) {
    ToolResult result;
    result.content.push_back(std::make_shared<types::TextContent>(text));
    return result;
}

ToolResult create_error_result(const std::string& error_message) {
    ToolResult result;
    result.is_error = true;
    result.content.push_back(std::make_shared<types::TextContent>(error_message));
    return result;
}

ToolResult create_image_result(const std::string& image_data, const std::string& mime_type) {
    ToolResult result;
    result.content.push_back(std::make_shared<types::ImageContent>(image_data, mime_type));
    return result;
}

ToolResult create_mixed_result(const std::vector<std::shared_ptr<types::ContentBlock>>& content_blocks) {
    ToolResult result;
    result.content = content_blocks;
    return result;
}

// Tool validation helpers
bool validate_tool_name(const std::string& name) {
    if (name.empty()) {
        return false;
    }

    // Tool names should be valid identifiers
    for (char c : name) {
        if (!(std::isalnum(c) || c == '_' || c == '-' || c == '.')) {
            return false;
        }
    }

    // Cannot start with a number or special character
    if (std::isdigit(name[0]) || name[0] == '-' || name[0] == '.') {
        return false;
    }

    return true;
}

std::string sanitize_tool_name(const std::string& name) {
    std::string result;
    result.reserve(name.size());

    for (char c : name) {
        if (std::isalnum(c) || c == '_' || c == '-' || c == '.') {
            result += c;
        } else {
            result += '_';
        }
    }

    // Ensure it doesn't start with invalid characters
    while (!result.empty() && (std::isdigit(result[0]) || result[0] == '-' || result[0] == '.')) {
        result.erase(0, 1);
    }

    return result.empty() ? "unnamed_tool" : result;
}

// Tool metadata helpers
json::Value create_tool_metadata(const std::string& category,
                                  const std::vector<std::string>& tags,
                                  const std::optional<std::string>& description) {
    json::Value metadata = json::Value::object({
        {"category", category},
        {"tags", json::Value::array()}
    });

    for (const auto& tag : tags) {
        metadata["tags"].push_back(tag);
    }

    if (description) {
        metadata["description"] = *description;
    }

    return metadata;
}

// Tool result builders
ToolResultBuilder::ToolResultBuilder() = default;
ToolResultBuilder& ToolResultBuilder::add_text(const std::string& text) {
    result_.content.push_back(std::make_shared<types::TextContent>(text));
    return *this;
}

ToolResultBuilder& ToolResultBuilder::add_image(const std::string& data, const std::string& mime_type) {
    result_.content.push_back(std::make_shared<types::ImageContent>(data, mime_type));
    return *this;
}

ToolResultBuilder& ToolResultBuilder::set_error(bool is_error) {
    result_.is_error = is_error;
    return *this;
}

ToolResultBuilder& ToolResultBuilder::set_structured_content(const json::Value& content) {
    result_.structured_content = content;
    return *this;
}

ToolResult ToolResultBuilder::build() && {
    return std::move(result_);
}

// Tool execution context
class ToolExecutionContext {
public:
    explicit ToolExecutionContext(const std::string& tool_name, Context<>& context)
        : tool_name_(tool_name), context_(context) {
        start_time_ = std::chrono::steady_clock::now();
    }

    ~ToolExecutionContext() {
        auto end_time = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time_);
        context_.info(std::format("Tool '{}' completed in {}ms", tool_name_, duration.count()));
    }

    void log_progress(double progress, const std::optional<std::string>& message = {}) {
        context_.report_progress(progress, 1.0, message);
    }

private:
    std::string tool_name_;
    Context<>& context_;
    std::chrono::steady_clock::time_point start_time_;
};

// Tool wrapper with timing and logging
class TimedToolWrapper {
public:
    static Result<ToolResult> execute(
        const std::string& tool_name,
        const std::function<Result<ToolResult>(const json::Value&, Context<>&)>& handler,
        const json::Value& arguments,
        Context<>& context
    ) {
        ToolExecutionContext exec_context(tool_name, context);

        context.info(std::format("Executing tool '{}'", tool_name));

        auto result = handler(arguments, context);

        if (result) {
            context.info(std::format("Tool '{}' completed successfully", tool_name));
        } else {
            context.error(std::format("Tool '{}' failed: {}", tool_name, result.error().message()));
        }

        return result;
    }
};

} // namespace server
} // namespace mcp